# -- coding: utf-8 --
"""
建立从训练集到tensorflow的文件流，全测试集模式
主要函数：
	inputPipeLine()
仅用作测试：
	show_distort()

"""   
import os
import tensorflow.python.platform
from six.moves import xrange 
import tensorflow as tf
from tensorflow.python.platform import gfile
import numpy as np
import pdb
import random
import cv2
from PIL import Image, ImageDraw
import define
from tensorflow.python.ops import control_flow_ops

#各网络默认输入size
img_resize = {'InceptionV3':299,
				'resnet_v2_50':224,
				'resnet_v2_101':224,
				'resnet_v2_152':224,
				'vgg':224,
				'nasnat':331,
				"InceptionResnetV2":299
				}

FLAGS = tf.app.flags.FLAGS

def readData(labelfile_queue):
	'''
	解析文件队列
	'''
	class ReadData(object):
		pass
	result = ReadData()
	reader = tf.TFRecordReader()
	_, serialized_example = reader.read(labelfile_queue) #key：行数；value：文件名
	features = tf.parse_single_example(serialized_example,
									   features={
											'image_raw': tf.FixedLenFeature([], tf.string),
											'label' : tf.FixedLenFeature([], tf.int64),
											'Xcoor' : tf.FixedLenFeature([], tf.int64),
											'Ycoor' : tf.FixedLenFeature([], tf.int64),
											'patchName' : tf.FixedLenFeature([], tf.string),
									   })
	img = tf.decode_raw(features['image_raw'], tf.uint8)
	img = tf.reshape(img, [256, 256, 3])
	result.image = img
	result.label = tf.cast(features['label'], tf.int32)
	result.Xcoor = tf.cast(features['Xcoor'], tf.int32)
	result.Ycoor = tf.cast(features['Ycoor'], tf.int32)
	result.filename = tf.cast(features['patchName'], tf.string)
	return result 		#TumorPatchFileName, label, Xcoor, Ycoor

def apply_with_random_selector(x, func, num_cases):
	sel = tf.random_uniform([], maxval=num_cases, dtype=tf.int32)
	# Pass the real x only to one of the func calls.
	return control_flow_ops.merge([
		func(control_flow_ops.switch(x, tf.equal(sel, case))[1], case)
		for case in range(num_cases)])[0]

def mean_image_subtraction(image):
	'''
	图像预处理：减去固定值或均值，对训练影响较大，注意调整，此处没用
	'''
	if image.get_shape().ndims != 3:
		raise ValueError('Input must be of size [height, width, C>0]')
	num_channels = image.get_shape().as_list()[-1]
	# if len(means) != num_channels:
	# 	raise ValueError('len(means) must match the number of channels')
	channels = tf.split(axis=2, num_or_size_splits=num_channels, value=image)
	for i in range(num_channels):
		# means = tf.reduce_mean(image[i])
		means=[123.68, 116.78, 103.94]
		channels[i] -= means[i]
	return tf.concat(axis=2, values=channels)

def distort_color(image, color_ordering=0, fast_mode=True, scope=None):
	'''
	图像色彩抖动，最后输出时加上tf.image.per_image_standardization(image)，注意这一步为必须，否则将严重影响训练结果

	Args：
	color_ordering：随机数，控制不同的参数抖动顺序
	fast_mode：是否启用快速处理模式，默认为是
	'''
	with tf.name_scope(scope, 'distort_color', [image]):
		if fast_mode:
			if color_ordering == 0:
				image = tf.image.random_brightness(image, max_delta=32. / 255.)
				image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
			else:
				image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
				image = tf.image.random_brightness(image, max_delta=32. / 255.)
		else:
			if color_ordering == 0:
				image = tf.image.random_brightness(image, max_delta=32. / 255.)
				image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
				image = tf.image.random_hue(image, max_delta=0.2)
				image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
			elif color_ordering == 1:
				image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
				image = tf.image.random_brightness(image, max_delta=32. / 255.)
				image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
				image = tf.image.random_hue(image, max_delta=0.2)
			elif color_ordering == 2:
				image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
				image = tf.image.random_hue(image, max_delta=0.2)
				image = tf.image.random_brightness(image, max_delta=32. / 255.)
				image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
			elif color_ordering == 3:
				image = tf.image.random_hue(image, max_delta=0.2)
				image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
				image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
				image = tf.image.random_brightness(image, max_delta=32. / 255.)
			else:
				raise ValueError('color_ordering must be in [0, 3]')
		return tf.image.per_image_standardization(image)

def inputPipeLine(is_train_data, batchSize = 32, fast_mode = True, Data_Dir = None, numEpochs = None, net_use=None, Data_path = None):
	'''
	Args:
	is_train_data：预处理数据集类型，仅对训练集做数据增强，测试集仅做提取
	net_use：所使用的网络，用于绑定默认输入图像尺寸
	num_epochs：指定每个文件读取多少遍（限制读取次数）？如果不指定，将无穷循环读取
	'''

	#创建文件队列
	# pdb.set_trace()
	if Data_path:
		#全测试集模式
		if not gfile.Exists(Data_path):
			raise ValueError('Failed to find Data_path: ' + Data_path)
		labelfile_queue = tf.train.string_input_producer([Data_path], shuffle=False, num_epochs = 1, capacity=10)
	elif Data_Dir:
		#单例测试模式，不打混，每个样本仅测试一次不补全batch
		if not gfile.Exists(Data_Dir):
			raise ValueError('Failed to find Data_Dir: ' + Data_Dir)
		filename_list =  []
		for filename in os.listdir(Data_Dir):
			if filename.endswith('.tfrecords'):
				filepath = os.path.join(Data_Dir,filename)
				filename_list.append(filepath)
		labelfile_queue = tf.train.string_input_producer(filename_list, shuffle=True, num_epochs = numEpochs) 
	else:
		raise ValueError('Data_Dir/Data_path is not specified')
	result = readData(labelfile_queue) #先读取,返回的是一条数据
	tf.summary.image('raw_image',
							tf.expand_dims(result.image, 0))
	image = result.image
	label = result.label
	
	if is_train_data == 'train':	#仅对训练数据做数据增强
		num_cases = 1 if fast_mode else 4
		image = apply_with_random_selector(
			image,
			lambda x, method: tf.image.resize_images(x, [img_resize[net_use],img_resize[net_use]], method=method),
			num_cases=num_cases)

		# image = tf.image.resize_images(image, [img_resize['resnet'],img_resize['resnet']])
		# tf.summary.image('resized_image',
		# 				 tf.expand_dims(image, 0))

		image = tf.image.random_flip_left_right(image)
		image = tf.image.random_flip_up_down(image)

		# 4种随机颜色都抖动方式
		image = apply_with_random_selector(
			image,
			lambda x, ordering: distort_color(x, ordering, fast_mode),
			num_cases=num_cases)

		if label==0:
			tf.summary.image('distorted_image/label-0',
				tf.expand_dims(image, 0))
		else:
			tf.summary.image('distorted_image/label-1',
				tf.expand_dims(image, 0))
		# image = tf.image.per_image_standardization(image)
	else:
		image = tf.image.resize_images(image, [img_resize[net_use],img_resize[net_use]])
		image = tf.image.per_image_standardization(image)

	if image.dtype != tf.float32:
		image = tf.image.convert_image_dtype(image, dtype=tf.float32)
	# tf.summary.image('final_image',tf.expand_dims(image, 0))
	if is_train_data == 'train':
		image_batch, label_batch, Xcoor_batch, Ycoor_batch, filename_batch = tf.train.shuffle_batch(
			[image, result.label, result.Xcoor, result.Ycoor, result.filename],
			batch_size = batchSize,
			num_threads=8,
			capacity=  300 * batchSize,
			allow_smaller_final_batch=False,
			min_after_dequeue= 200 * batchSize)
		return image_batch, label_batch, Xcoor_batch, Ycoor_batch, filename_batch
	else:
		#此处增加rawimage_batch，仅用于后期展示网络输入的原图batch用
		image_batch, label_batch, Xcoor_batch, Ycoor_batch, filename_batch, rawimage_batch = tf.train.shuffle_batch(
			[image, result.label, result.Xcoor, result.Ycoor, result.filename, result.image],
			batch_size = batchSize,
			num_threads=8,
			capacity=  300 * batchSize,
			allow_smaller_final_batch=True,
			min_after_dequeue= 200 * batchSize)
	
		return image_batch, label_batch, Xcoor_batch, Ycoor_batch, filename_batch, rawimage_batch 

def show_distort(image_batch, sess):
	'''
	仅用于演示颜色抖动效果

	'''
	image_raw = tf.squeeze(image_batch)
	image_standar = tf.image.per_image_standardization(image_raw)
	image_bright = tf.image.random_brightness(image_raw, max_delta=32. / 255.)
	image_sat = tf.image.random_saturation(image_raw, lower=0.5, upper=1.5)
	image_hue = tf.image.random_hue(image_raw, max_delta=0.2)
	image_cont = tf.image.random_contrast(image_raw, lower=0.5, upper=1.5)
	image_raw, image_standar, image_bright, image_sat, image_hue, image_cont = sess.run([image_raw, image_standar, image_bright, image_sat, image_hue, image_cont])

	patch_save_path = define.REFERANCE_IMG_SAVE_DIR + '/distort_img'

	# cv2.imshow('raw',np.squeeze(np.array(image_raw.tolist())))
	# cv2.waitKey(0)
	patch_save_path_1 = patch_save_path + '/image_raw.jpg'
	cv2.imwrite(patch_save_path_1, np.array(image_raw.tolist())*255)

	# cv2.imshow('per_image_standardization',np.squeeze(np.array(image_standar.tolist())))
	# cv2.waitKey(0)
	patch_save_path_2 = patch_save_path + '/image_standar.jpg'
	cv2.imwrite(patch_save_path_2, np.array(image_standar.tolist())*255)

	# cv2.imshow('random_brightness',np.squeeze(np.array(image_bright.tolist())))
	# cv2.waitKey(0)
	patch_save_path_3 = patch_save_path + '/image_bright.jpg'
	cv2.imwrite(patch_save_path_3, np.array(image_bright.tolist())*255)

	# cv2.imshow('random_saturation',np.squeeze(np.array(image_sat.tolist())))
	# cv2.waitKey(0)
	patch_save_path_4 = patch_save_path + '/image_sat.jpg'
	cv2.imwrite(patch_save_path_4, np.array(image_sat.tolist())*255)

	# cv2.imshow('random_hue',np.squeeze(np.array(image_hue.tolist()))*255)
	# cv2.waitKey(0)
	patch_save_path_5 = patch_save_path + '/image_hue.jpg'
	cv2.imwrite(patch_save_path_5, np.array(image_hue.tolist())*255)

	# cv2.imshow('random_contrast',np.squeeze(np.array(image_cont.tolist())))
	# cv2.waitKey(0)
	patch_save_path_6 = patch_save_path + '/image_cont.jpg'
	cv2.imwrite(patch_save_path_6, np.array(image_cont.tolist())*255)

def main():
	'''
	测试函数，根据需要测试函数

	'''
	#数据源路径自定义
	DATASET_DIR='/home/ramsley/DataSet/prostate/trainingData/patch-based-classification/tf-records-balance'
	is_train_data = 'validation'
	total_neg=0
	total_pos=0
	total_examples=0
	image_batch, label_batch, Xcoor_batch, Ycoor_batch, filename_batch = inputPipeLine(is_train_data, batchSize = 1, \
																fast_mode = True, Data_Dir = DATASET_DIR, numEpochs = 1)
	total_files = 0
	init_op = tf.initialize_local_variables()	#是否需要
	# with tf.Graph().as_default(), tf.device('/cpu:0'):
	gpu_options = tf.GPUOptions(allow_growth=True)
	# init_op = tf.global_variables_initializer()
	loc_init_op = tf.local_variables_initializer()
	sess = tf.Session(config=tf.ConfigProto(
			allow_soft_placement=True,
			gpu_options=gpu_options))
	sess.run([init_op,loc_init_op])

	# sess.run(init_op)   # 初始化本地变量 	
	coord = tf.train.Coordinator()
	threads = tf.train.start_queue_runners(sess=sess, coord=coord)
	try:
		while not coord.should_stop():
			label = label_batch.eval(session=sess)
			# show_distort(image_batch, sess)
			# pdb.set_trace()
			# if label==0:
			# 	total_neg += 1
			# else:
			# 	total_pos += 1
			total_examples += 1
			# if total_examples%1000==0:
				# file = filename.eval(session=sess)
			# print(total_examples)
			print(label)
			pdb.set_trace()
	except tf.errors.OutOfRangeError:
		print('Epochs Complete!')
	finally:
		coord.request_stop()
	coord.join(threads)
	sess.close()
	print total_files

if __name__ == '__main__':
	main()